FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

# 设置工作目录
WORKDIR /app

# 设置环境变量
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# 为 apt 设置清华源
RUN sed -i 's@//.*archive.ubuntu.com@//mirrors.tuna.tsinghua.edu.cn@g' /etc/apt/sources.list && \
    sed -i 's@//.*security.ubuntu.com@//mirrors.tuna.tsinghua.edu.cn@g' /etc/apt/sources.list

# 更新包列表并安装系统依赖（包括编译工具）
RUN apt update && \
    apt install -y --no-install-recommends \
    git \
    wget \
    curl \
    ca-certificates \
    build-essential \
    cmake \
    libc6-dev \
    && rm -rf /var/lib/apt/lists/*

# 设置编译器环境变量
ENV CC=/usr/bin/gcc
ENV CXX=/usr/bin/g++

# 为 pip 设置清华源并安装Python依赖
RUN pip install --no-cache-dir --upgrade pip -i https://pypi.tuna.tsinghua.edu.cn/simple && \
    pip install --no-cache-dir \
    transformers==4.34.0 \
    accelerate \
    sentencepiece \
    protobuf \
    torchvision \
    datasets \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 llama-cpp-python（现在应该可以从源码成功编译了）
RUN pip install --no-cache-dir llama-cpp-python \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

# 创建模型目录
RUN mkdir -p /models

# 设置默认的启动命令
CMD ["bash"]